import pandas as pd
from typing import List

from interfaces.transformers import IBaseSubjects
from interfaces.data_loader import IDataLoader
from interfaces.form_metadata import IFormMetadataProvider
from constants import (
    ICF_COLS, SEX_COLS, SUBJECT_NUMBER, DATA_FESTIVAL,LINE_NUMBER,
    ECSTDAT, ECSTTIM, DRUG_COLS, ENROLLMENT_COLS, ENROLLMENT_COLS_DSFAST,
    FormAllConstants,
)


class BaseSubjectsTransformer(IBaseSubjects):
    def __init__(self, data_loader: IDataLoader, metadata: IFormMetadataProvider):
        self.data_loader = data_loader
        self.metadata = metadata

    def get_qualified_subjects(self) -> pd.DataFrame:
        result_ICF_ModuleOID = self.metadata.get_targets((FormAllConstants.IsICF, True), FormAllConstants.ModuleOID, True)
        result_IsDemographics_ModuleOID = self.metadata.get_targets((FormAllConstants.IsDemographics, True), FormAllConstants.ModuleOID, True)
        result_Enrollment_ModuleOID = self.metadata.get_targets((FormAllConstants.IsEnrollment, True), FormAllConstants.ModuleOID, True)
        result_Summary_ModuleOID = self.metadata.get_targets((FormAllConstants.IsSummary, True), FormAllConstants.ModuleOID, True)
        result_Drug_ModuleOID = self.metadata.get_targets((FormAllConstants.IsDrug, True), FormAllConstants.ModuleOID, True)
        result_Trial_FolderName = self.metadata.get_targets((FormAllConstants.IsTrial, True), FormAllConstants.FolderName, True)

        df_ICF = self.data_loader.load_specific(result_ICF_ModuleOID[0], ICF_COLS)
        df_SEX = self.data_loader.load_specific(result_IsDemographics_ModuleOID[0], SEX_COLS)
        df_Enrollment = self.data_loader.load_specific(result_Enrollment_ModuleOID[0], ENROLLMENT_COLS)
        # DSFAST 可选加载，失败则跳过
        try:
            df_Enrollment_DSFAST = self.data_loader.load_specific(result_Enrollment_ModuleOID[0], ENROLLMENT_COLS_DSFAST)
        except Exception:
            df_Enrollment_DSFAST = pd.DataFrame(columns=[SUBJECT_NUMBER, 'DSFAST'])
        df_Summary = self.data_loader.load_sheet(result_Summary_ModuleOID[0])

        df_DRUG = self.data_loader.load_specific(result_Drug_ModuleOID[0], DRUG_COLS)
        df_DRUG_filtered = df_DRUG[df_DRUG[DATA_FESTIVAL].isin(result_Trial_FolderName)].copy()

        ec_pivoted = df_DRUG_filtered.pivot_table(
            index=SUBJECT_NUMBER,
            columns=[DATA_FESTIVAL, LINE_NUMBER],
            values=[ECSTDAT, ECSTTIM],
            aggfunc='first'
        )
        row_numbers = sorted(df_DRUG_filtered[LINE_NUMBER].unique().tolist()) if LINE_NUMBER in df_DRUG_filtered.columns else []
        new_columns: List = []
        for period in result_Trial_FolderName:
            for row_num in row_numbers:
                if (ECSTDAT, period, row_num) in ec_pivoted.columns:
                    new_columns.append((ECSTDAT, period, row_num))
                if (ECSTTIM, period, row_num) in ec_pivoted.columns:
                    new_columns.append((ECSTTIM, period, row_num))
        ec_pivoted_filtered = ec_pivoted[new_columns].copy()
        ec_pivoted_filtered.columns = [f"{period}_{row_num}_{field}" for field, period, row_num in ec_pivoted_filtered.columns]

        入组成功受试者列表 = (
            df_ICF
            .merge(df_SEX, on=SUBJECT_NUMBER, how='left')
            .merge(df_Enrollment, on=SUBJECT_NUMBER, how='left')
            .merge(ec_pivoted_filtered, on=SUBJECT_NUMBER, how='left')
            .merge(df_Summary, on=SUBJECT_NUMBER, how='left')
        )
        # 如存在 DSFAST，则合并并将其置于前列
        if not df_Enrollment_DSFAST.empty and 'DSFAST' in df_Enrollment_DSFAST.columns:
            入组成功受试者列表 = 入组成功受试者列表.merge(df_Enrollment_DSFAST, on=SUBJECT_NUMBER, how='left')
            cols = list(入组成功受试者列表.columns)
            if 'DSFAST' in cols:
                reordered = ['DSFAST', SUBJECT_NUMBER] + [c for c in cols if c not in ['DSFAST', SUBJECT_NUMBER]]
                入组成功受试者列表 = 入组成功受试者列表[reordered]
        return 入组成功受试者列表
